#include "ShadowRenderPass.h"

#include <SERender/Renderer/Renderer.h>
#include <SERender/Descriptors/FrameDescriptor.h>
#include <SEngine/Render/Descriptors/SceneDescriptor.h>
#include <SECore/Ecs/Scene.h>
#include <SEngine/Ecs/Components/ModelRenderComp.h>
#include <SEngine/Render/RenderFeature/ModelRenderFeature.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>
#include <SEEdit/Render/RenderFeature/DebugRenderFeature.h>
#include <SEngine/Ecs/Components/LightComp.h>
#include <SERender/Resourses/UniformBuffer.h>
#include <SEngine/Ecs/Components/TransformComp.h>
#include <SERender/Resourses/Loader/TextureLoader.h>
#include <SEngine/Render/RenderFeature/UBORenderFeature.h>
#include <SEngine/Render/RenderFeature/UtilRenderFeature.h>
#include <SEngine/Ecs/Components/UuidComp.h>
#include <SECore/Manager.h>

#include <SERender/Resourses/Loader/VertexArrayLoader.h>
#include <SERender/Resourses/Loader/ShaderLoader.h>
#include <SEUi/SEUi.h>

namespace SEngine
{


    ShadowRenderPass::ShadowRenderPass(Renderer &renderer)
        : RenderPass(renderer) 
    {
        {
            m_fbo = std::make_unique<FrameBuffer>();

            TextureLoader<Texture2D> texLoader;
            texLoader.SetInternalFormat(RGB32F);
            texLoader.SetFormat(RGB);
            texLoader.SetType(FLOAT);
            texLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
            texLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
            texLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
            texLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);
            m_positionRT = texLoader.CreateTexture(1024, 1024, nullptr);
            m_normalRT = texLoader.CreateTexture(1024, 1024, nullptr);
            m_shadowMap = texLoader.CreateTexture(1024, 1024, nullptr);
            texLoader.SetInternalFormat(RGBA);
            texLoader.SetFormat(RGBA);
            texLoader.SetType(UNSIGNED_BYTE);
            m_objectIDMap = texLoader.CreateTexture(1024, 1024, nullptr);
            m_lightViewObjectIDMap = texLoader.CreateTexture(1024, 1024, nullptr);

            // 线光源阴影贴图
            texLoader.SetInternalFormat(RGB32F);
            texLoader.SetFormat(RGB);
            texLoader.SetType(FLOAT);
            m_directLightDepthTex = texLoader.CreateTexture(1024, 1024, nullptr);
            
            // 点光源阴影贴图
            TextureLoader<TextureCube> cubeTexLoader;
            cubeTexLoader.SetInternalFormat(DEPTH_COMPONENT);
            cubeTexLoader.SetFormat(DEPTH_COMPONENT);
            cubeTexLoader.SetType(UNSIGNED_BYTE);
            cubeTexLoader.SetParameter(TEXTURE_MAG_FILTER, LINEAR);
            cubeTexLoader.SetParameter(TEXTURE_MIN_FILTER, LINEAR);
            cubeTexLoader.SetParameter(TEXTURE_WRAP_S, CLAMP_TO_EDGE);
            cubeTexLoader.SetParameter(TEXTURE_WRAP_T, CLAMP_TO_EDGE);
            cubeTexLoader.SetParameter(TEXTURE_WRAP_R, CLAMP_TO_EDGE);
            m_pointLightDepthTex = cubeTexLoader.CreateTexture(1024, 1024, nullptr);
            // 点光源阴影帧缓冲
            m_pointLightDepthFbo = std::make_unique<FrameBuffer>();
            m_pointLightDepthFbo->SwapAttachments(Attachment::DEPTH_ATTACHMENT, *m_pointLightDepthTex);
            m_pointLightDepthFbo->EnableRenderBuffer(false);


        }   

        // // 深度贴图
        // texLoader.SetInternalFormat(RGB32F);
        // texLoader.SetFormat(RGB);
        // texLoader.SetType(FLOAT);
        // texLoader.SetParameter(TEXTURE_MIN_FILTER, NEAREST);
        // texLoader.SetParameter(TEXTURE_MAG_FILTER, NEAREST);
        // fboLoader.AddAttachment(texLoader.CreateTexture(800, 600, nullptr), Attachment::COLOR_ATTACHMENT0);
        // fboLoader.SetRenderBuffer(true);
        // m_depthFbo = fboLoader.CreateFrameBuffer(800, 600);

        // 阴影渲染对象
        m_directLightShadow = Manager::GetInstance().Get<Material>("DirectLightShadow").value();
        m_directLightShadow.Value().GetPSO().blend = true;
        m_directLightShadow.Value().GetPSO().depth = false;
        
        m_cubeDepthShadowMaterial = Manager::GetInstance().Get<Material>("PointLightShadow").value();
        m_cubeDepthShadowMaterial.Value().GetPSO().blend = true;
        m_cubeDepthShadowMaterial.Value().GetPSO().depth = false;
    

        // 渲染单张深度图
        m_depthMaterial = Manager::GetInstance().Get<Material>("Depth").value();
        m_depthMaterial.Value().GetPSO().depth = true;

        // 渲染Cube深度图(6张)
        m_cubeDepthMaterial = Manager::GetInstance().Get<Material>("CubeDepth").value();
        m_cubeDepthMaterial.Value().GetPSO().depth = true;
    }

    void ShadowRenderPass::BeginFrame()
    {
        auto & sceneDes = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get();
        auto & frameDes = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & modelFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        // 初始化gBuffer
        // 渲染相机角度的深度图
        m_renderer.UseFrameBuffer(*m_fbo, [&]() {
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            m_fbo->ClearAttachments();
            m_fbo->EnableRenderBuffer(true);
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT0, *m_positionRT);
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT1, *m_objectIDMap);
            m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT2, *m_normalRT);
            m_fbo->Resize(frameDes.w, frameDes.h);
            m_fbo->Use();
            m_renderer.GetDrive()->ClearColor(0.0f, 0.0f, 0.0f, 0.0f);
            m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
            sceneDes.scene->Components<ModelRenderComp, UuidComp>([&](Actor & actor, ModelRenderComp & modelComp, UuidComp & uuidComp) {
                uint32_t id = uuidComp.Hash<uint32_t>();
                m_depthMaterial.Value().Set("u_id", glm::vec4{
                    ((uint8_t*)&id)[0] / 255.0f,
                    ((uint8_t*)&id)[1] / 255.0f,
                    ((uint8_t*)&id)[2] / 255.0f,
                    ((uint8_t*)&id)[3] / 255.0f
                });
                modelFeature.DrawModel(modelComp.m_model, m_depthMaterial);
            });
        });
        // m_cubeDepthShadowMaterial->Set("u_cameraDepthMap", Shader::Sampler{m_positionRT, TextureSlot::TEXTURE0});
        
        // ImGui::Begin("RT");
        // if (ImGui::Button("Reload")) {
        //     m_directLightShadow->SetShader(ShaderLoader().Load("D:\\WorkSpace\\serious-engine\\Engine\\Shader\\DirectLightShadow.glsl").CreateShader());
        // }
        // ImGui::Image((ImTextureID)m_normalRT->GetId(), ImVec2{(float)frameDes.w, (float)frameDes.h}, ImVec2{0, 1}, ImVec2{1, 0});
        // ImGui::End();
    }
    
    struct LightCameraInfo
    {
        Camera camera;
        float zNear;
        float zFar;
    };
    LightCameraInfo CalCamera(DirectLight & light, int w, int h, float zNear, float zFar, Camera viewCamera) {
        viewCamera.SetZNear(zNear);
        viewCamera.SetZFar(zFar);
        auto inpv = glm::inverse(viewCamera.GetProjectMatrix(w, h) * viewCamera.GetViewMatrix());
        auto GetWorldPos = [inpv] (glm::vec3 vertex) {
            auto v = inpv * glm::vec4{vertex, 1.0f};
            return glm::vec3(v / v.w);
        };
        std::vector<glm::vec3> frustum{
            GetWorldPos(glm::vec3{-1.0f, -1.0f, -1.0f}),
            GetWorldPos(glm::vec3{1.0f, -1.0f, -1.0f}),
            GetWorldPos(glm::vec3{1.0f, 1.0f, -1.0f}),
            GetWorldPos(glm::vec3{-1.0f, 1.0f, -1.0f}),
            GetWorldPos(glm::vec3{-1.0f, -1.0f, 1.0f}),
            GetWorldPos(glm::vec3{1.0f, -1.0f, 1.0f}),
            GetWorldPos(glm::vec3{1.0f, 1.0f, 1.0f}),
            GetWorldPos(glm::vec3{-1.0f, 1.0f, 1.0f}),
        };
        glm::vec3 center{0.0f};
        for (int i = 0; i < 8; ++i) {
            center += (frustum[i] / 8.0f);
        }
        float radius = -1.0f;
        for (int i = 0; i < 8; ++i) {
            radius = glm::max(radius, glm::distance(center, frustum[i]));
        }
        const float rangeEx = radius * 5;
        Camera camera;
        camera.SetUp(glm::vec3{0.0f, 1.0f, 0.0f});
        camera.SetPos(glm::vec3(center) + glm::normalize(-light.GetDirect()) * (radius + rangeEx));
        camera.SetFront(center - camera.GetPos());
        camera.SetZNear(0.1f);
        camera.SetZFar((radius * 2 + rangeEx));
        camera.SetSize(radius);
        camera.SetType(CameraType::Ortho);
        return LightCameraInfo{camera, zNear, zFar};
    }

    void ShadowRenderPass::DrawFrame()
    {

        CreateDirectLightShadowMap();
        CreatePointLightShadowMap();
    }   

    void ShadowRenderPass::CreateDirectLightShadowMap()
    {
        auto & frameDes = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & sceneDes = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get();
        auto & uboFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UBORenderFeature>()->get();
        auto & modelFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        auto & debugFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<DebugRenderFeature>()->get();
        auto & utilRenderFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UtilRenderFeature>()->get();
        
        ImGui::Begin("RT");
        static float bias = 0.005f;
        ImGui::DragFloat("Bias", &bias, 0.01f, 0.0f, 0.05f);
        ImGui::End();

        // 渲染平行光阴影图
        sceneDes.scene->Components<DirectLightComp, TransformComp>([&](Actor & actor, DirectLightComp & dirLight, TransformComp & transComp) {
            auto & light = dirLight.light;
            
            float zNear = frameDes.camera.GetZNear();
            float zFar = glm::min(frameDes.camera.GetZFar(), 200.0f);
            auto ZValue = [=](float i, float N, float lambda = 0.7) -> float {
                return lambda * zNear * glm::pow(zFar / zNear, i / N) + (1 - lambda) * glm::mix(zNear, zFar, i / N);
            };
            LightCameraInfo cameraNear = CalCamera(light, frameDes.w, frameDes.h, ZValue(0, 3), ZValue(1, 3), frameDes.camera);
            LightCameraInfo cameraMid = CalCamera(light, frameDes.w, frameDes.h, ZValue(1, 3), ZValue(2, 3), frameDes.camera);
            LightCameraInfo cameraFar = CalCamera(light, frameDes.w, frameDes.h, ZValue(2, 3), ZValue(3, 3), frameDes.camera);

            std::vector<LightCameraInfo> cameras{cameraNear, cameraMid, cameraFar};
            
            LightCameraInfo cameraAll = CalCamera(light, frameDes.w, frameDes.h, frameDes.camera.GetZNear(), frameDes.camera.GetZFar(), frameDes.camera);
            
            m_directLightShadow.Value().Set("u_camPos", frameDes.camera.GetPos());
            m_directLightShadow.Value().Set("u_camFront", frameDes.camera.GetFront());

            auto DrawShadow = [&](Camera & camera, float zNear, float zFar) {
                m_renderer.UseFrameBuffer(*m_fbo, [&]() {
                    uboFeature.UpdateCameraUBOData(camera, 2048, 2048);
                    m_renderer.GetDrive()->Viewport(0, 0, 2048, 2048);
                    m_fbo->ClearAttachments();
                    m_fbo->EnableRenderBuffer(true);
                    m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT1, *m_lightViewObjectIDMap);
                    m_fbo->SwapAttachments(Attachment::COLOR_ATTACHMENT3, *m_directLightDepthTex);
                    m_fbo->Resize(2048, 2048);
                    m_fbo->Use();
                    m_renderer.GetDrive()->Clear(COLOR_BUFFER_BIT | DEPTH_BUFFER_BIT);
                    sceneDes.scene->Components<ModelRenderComp, UuidComp>([&](Actor & actor, ModelRenderComp & modelComp, UuidComp & uuidComp) {
                        uint32_t id = uuidComp.Hash<uint32_t>();
                        m_depthMaterial.Value().Set("u_id", glm::vec4{
                            ((uint8_t*)&id)[0] / 255.0f,
                            ((uint8_t*)&id)[1] / 255.0f,
                            ((uint8_t*)&id)[2] / 255.0f,
                            ((uint8_t*)&id)[3] / 255.0f
                        });
                        modelFeature.DrawModel(modelComp.m_model, m_depthMaterial);
                    });
                    m_directLightShadow.Value().Set("u_objectIDMap", Shader::Sampler{*m_objectIDMap, TextureSlot::TEXTURE0});
                    m_directLightShadow.Value().Set("u_positionRT", Shader::Sampler{*m_positionRT, TextureSlot::TEXTURE1});
                    m_directLightShadow.Value().Set("u_normalRT", Shader::Sampler{*m_normalRT, TextureSlot::TEXTURE2});
                    m_directLightShadow.Value().Set("u_lightDepthMap", Shader::Sampler{*m_directLightDepthTex, TextureSlot::TEXTURE3});
                    m_directLightShadow.Value().Set("u_lightViewObjectIDMap", Shader::Sampler{*m_lightViewObjectIDMap, TextureSlot::TEXTURE4});
                    m_directLightShadow.Value().Set("u_lightDir", light.GetDirect());
                    m_directLightShadow.Value().Set("u_lightPV", camera.GetProjectMatrix(2048, 2048) * camera.GetViewMatrix());
                    m_directLightShadow.Value().Set("u_zNear", zNear);
                    m_directLightShadow.Value().Set("u_zFar", zFar);
                    m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
                    uboFeature.UpdateCameraUBOData(frameDes.camera, frameDes.w, frameDes.h);
                });
                utilRenderFeature.DrawBillboard(m_directLightShadow);
            };
            for (auto & cameraInfo : cameras) {
                DrawShadow(cameraInfo.camera, cameraInfo.zNear, cameraInfo.zFar);
            }
            // DrawShadow(cameraAll.camera, cameraAll.zNear, cameraAll.zFar);
        });
    }

    void ShadowRenderPass::CreatePointLightShadowMap()
    {
        auto & frameDes = m_renderer.AnyDescriptor::GetDescriptor<FrameDescriptor>()->get();
        auto & sceneDes = m_renderer.AnyDescriptor::GetDescriptor<SceneDescriptor>()->get();
        auto & uboFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UBORenderFeature>()->get();
        auto & modelFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<ModelRenderFeature>()->get();
        auto & debugFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<DebugRenderFeature>()->get();
        auto & utilRenderFeature = m_renderer.Descriptor<RenderFeature>::GetDescriptor<UtilRenderFeature>()->get();
        
        // 渲染点光源阴影图
        // m_shadowBoard.SetMaterial(m_cubeDepthShadowMaterial);
        sceneDes.scene->Components<PointLightComp, TransformComp>([&](Actor & actor, PointLightComp & ptLight, TransformComp & transComp) {
            m_renderer.GetDrive()->Viewport(0, 0, 1024, 1024); 
            auto & light = ptLight.light;
            std::vector<glm::vec3> fronts{
                glm::vec3{1.0f, 0.0f, 0.0f},
                glm::vec3{-1.0f, 0.0f, 0.0f},
                glm::vec3{0.0f, 1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, 0.0f, 1.0f},
                glm::vec3{0.0f, 0.0f, -1.0f},
            };
            std::vector<glm::vec3> ups{
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, 0.0f, 1.0f},
                glm::vec3{0.0f, 0.0f, -1.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},
                glm::vec3{0.0f, -1.0f, 0.0f},

            };
            const float zFar = 500.0f;
            for (int i = 0; i < 6; ++i) {
                Camera camera;
                camera.SetFov(90.0f);
                camera.SetPos(light.GetPos());
                camera.SetZNear(0.1f);
                camera.SetZFar(zFar);
                camera.SetUp(ups[i]);
                camera.SetFront(fronts[i]);
                auto vp = camera.GetProjectMatrix(1024, 1024) * camera.GetViewMatrix();
                m_cubeDepthMaterial.Value().Set("u_shadowMat[" + std::to_string(i) + "].matrix", vp);
            }
            m_cubeDepthMaterial.Value().Set("u_lightPos", light.GetPos());
            m_cubeDepthMaterial.Value().Set("u_zFar", zFar);
            m_renderer.UseFrameBuffer(*m_pointLightDepthFbo, [&]() {
                m_renderer.GetDrive()->Clear(DEPTH_BUFFER_BIT);
                sceneDes.scene->Components<ModelRenderComp>([&](Actor & actor, ModelRenderComp & modelComp) {
                    modelFeature.DrawModel(modelComp.m_model, m_cubeDepthMaterial);
                });
            });
            m_cubeDepthShadowMaterial.Value().Set("u_lightDepthMap", Shader::Sampler{*m_pointLightDepthTex, TextureSlot::TEXTURE1});
            m_cubeDepthShadowMaterial.Value().Set("u_lightPos", light.GetPos());
            m_cubeDepthShadowMaterial.Value().Set("u_zFar", zFar);
            m_cubeDepthShadowMaterial.Value().Set("u_bias", 0.05f);
            m_renderer.GetDrive()->Viewport(frameDes.x, frameDes.y, frameDes.w, frameDes.h);
            utilRenderFeature.DrawBillboard(m_cubeDepthShadowMaterial);
        });
    }


} // namespace SEngine
